﻿// Example 20-01. Using K-means
#include <iostream>
#include <opencv2/opencv.hpp>

using namespace std;
constexpr int MAX_CLUSTERS = 5;

static void   help(char* argv[])
{
    cout << "\nThis program demonstrates kmeans clustering.\n"
            " It generates an image with random points, then assigns a random number\n"
            " of cluster centers and uses kmeans to move those cluster centers to their\n"
            " representative location\n"
            "Usage:\n"
         << argv[0] << "\n\n"
         << "ESC or 'q' or 'Q' to quit\n\n"
         << endl;
}

int main(int argc, char** argv)
{
    help(argv);
    const cv::Scalar colorTab[] = { cv::Scalar(0, 0, 255),
                                    cv::Scalar(0, 255, 0),
                                    cv::Scalar(255, 100, 100),
                                    cv::Scalar(255, 0, 255),
                                    cv::Scalar(0, 255, 255) };
    cv::Mat          img(500, 500, CV_8UC3);
    cv::RNG          rng(time(nullptr));

    while (true) {
        int     clusterCount = rng.uniform(2, MAX_CLUSTERS + 1);
        int     sampleCount  = rng.uniform(100, 1001); // 提高最小样本数
        cv::Mat points(sampleCount, 1, CV_32FC2), labels;
        clusterCount = std::min(clusterCount, sampleCount);
        cv::Mat centers(clusterCount, 1, points.type());

        // 生成随机样本
        for (int k = 0; k < clusterCount; k++) {
            cv::Point center(rng.uniform(50, img.cols - 50), // 避免边缘
                             rng.uniform(50, img.rows - 50));
            cv::Mat   pointChunk
                = points.rowRange(k * sampleCount / clusterCount,
                                  k == clusterCount - 1 ? sampleCount : (k + 1) * sampleCount / clusterCount);
            rng.fill(pointChunk,
                     cv::RNG::NORMAL,
                     cv::Scalar(center.x, center.y),
                     cv::Scalar(img.cols * 0.05, img.rows * 0.05));
        }

        cv::randShuffle(points, 1, &rng);
        cv::kmeans(points,
                   clusterCount,
                   labels,
                   cv::TermCriteria(cv::TermCriteria::EPS | cv::TermCriteria::COUNT, 10, 1.0),
                   3,
                   cv::KMEANS_PP_CENTERS,
                   centers);

        img = cv::Scalar::all(0);
        // 绘制样本点
        for (int i = 0; i < sampleCount; i++) {
            int       clusterIdx = labels.at<int>(i);
            cv::Point ipt        = points.at<cv::Point2f>(i);
            cv::circle(img, ipt, 2, colorTab[clusterIdx], cv::FILLED, cv::LINE_AA);
        }

        // 绘制中心点
        for (int k = 0; k < centers.rows; k++) {
            // 更安全的访问方式
            cv::Mat center_mat;
            centers.row(k).convertTo(center_mat, CV_32F);  // 保持单通道
            if (auto* data = center_mat.ptr<float>(0); data && center_mat.total() >= 2) {  // 确保有足够数据
                cv::Point2f center(data[0], data[1]);
                cv::circle(img, center, 10, cv::Scalar(255, 255, 255), 2, cv::LINE_AA);
            }
        }

        cv::imshow("Example 20-01", img);
        if (auto key = cv::waitKey(); key == 27 || key == 'q' || key == 'Q') // 'ESC'
            break;
    }
    return 0;
}
